feat: add jemalloc heap profiling infrastructure by eudelins-zama · Pull Request #449 · zama-ai/kms

eudelins-zama · 2026-03-05T12:24:45Z

Description of changes

Add jemalloc heap profiling infrastructure to detect memory leaks in KMS core nodes.

Bumped trivy-action to avoid this CI failure.

What's included

Feature-gated jemalloc integration (heap-profiling Cargo feature):

Replaces the default allocator with tikv-jemallocator when the feature is enabled
Installs a SIGUSR1 signal handler that triggers on-demand heap dumps to /tmp/kms-heap/
Exposes two new Prometheus metrics (kms_jemalloc_allocated, kms_jemalloc_resident) via a jemalloc-stats feature on the observability crate, enabling leak-type diagnosis (application leak vs allocator fragmentation vs non-jemalloc growth)

Dedicated Cargo profile (heap-profiling):

Inherits from release with debug=1 (line tables only) and strip=none so jeprof can resolve addresses to function:line without the full DWARF overhead

Docker Compose overlay (profiling/docker-compose-heap-profiling.yml):

Overrides build args to enable the feature, profile, and -C force-frame-pointers=yes for reliable backtraces
Configures MALLOC_CONF with prof:true, lg_prof_sample:12, prof_gdump:true, and prof_final:true

Makefile targets:

build-compose-heap-profiling / start-compose-heap-profiling / stop-compose-heap-profiling
dump-heap-profiles: sends SIGUSR1 to all 4 cores, copies .heap files, the binary, and /proc/PID/maps locally

Analysis script (profiling/analyze-heap.sh):

Cross-platform (Linux + macOS) script that handles PIE/ASLR address resolution by rewriting MAPPED_LIBRARIES paths and injecting maps.txt. Not tested on macOS though.
Produces 4 outputs: top-leaks.txt, latest.svg, diff-leaks.txt (allocation sites that grew), and diff.svg (diff flamegraph)
Falls back to manual addr2line resolution when jeprof can't resolve symbols

Dockerfile changes:

Adds CARGO_EXTRA_FEATURES and RUSTFLAGS build args so the profiling stack can inject feature flags and frame pointers
Uses profile-keyed Docker cache IDs (id=cargo-target-${LTO_RELEASE}) to avoid cache collisions between release and heap-profiling builds

Config fix (separate commit):

Adds missing reshare field to all compose config files

Issue ticket number and link

Closes https://github.com/zama-ai/kms-internal/issues/2927

PR Checklist

I attest that all checked items are satisfied. Any deviation is clearly justified above.

Dependency Update Questionnaire (only if deps changed or added)

Answer in the Cargo.toml next to the dependency (or here if updating):

Ownership changes or suspicious concentration?
Low popularity?
Unusual version jump?
Lacking documentation?
Missing CI?
No security / disclosure policy?
Significant size increase?

More details and explanations for the checklist and dependency updates can be found in CONTRIBUTING.md

eudelins-zama · 2026-03-05T12:41:05Z

Usage

# Build and start env
make build-compose-heap-profiling
make start-compose-heap-profiling

# Keygen
cargo run --bin kms-core-client -- -f core-client/config/client_local_threshold.toml -l -a insecure-key-gen
# copy key id to be used below:

# First decryption burst to have representative memory usage after some use
cargo run --bin kms-core-client -- -f core-client/config/client_local_threshold.toml -l -a public-decrypt from-args --to-encrypt 0xC2 --data-type euint64 -b 1 -n 2000 --key-id $KEY_ID --inter-request-delay-ms 500 -p 40

# First memory dump
make dump-heap-profiles

# Second decryption burst
cargo run --bin kms-core-client -- -f core-client/config/client_local_threshold.toml -l -a public-decrypt from-args --to-encrypt 0xC2 --data-type euint64 -b 1 -n 2000 --key-id $KEY_ID --inter-request-delay-ms 500 -p 40

# Second memory dump
make dump-heap-profiles

# Compare memory  usage for core-1 before/after the second bump
./profiling/analyze-heap.sh ./profiling/heap-dumps/kms-server ./profiling/heap-dumps/core-1/

# Analysis of the `./profiling/heap-analysis/diff-leaks.txt` file -> Claude is pretty good at this

github-actions · 2026-03-05T13:19:36Z

Consolidated Tests Results 2026-03-08 - 20:15:43

Test Results

15 passed

Details

15 tests
not captured
junit-to-ctrf
build-and-test test-reporter #749
feat: add jemalloc heap profiling infrastructure #449

test-reporter: Run #749

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
15	15	0	0	0	0	0	not captured

🎉 All tests passed!

Tests

View All Tests

Test Name	Status	Duration
nightly_full_gen_tests_k8s_default_threshld_sequential_crs	✅	33.0s
test_k8s_threshld_insecure	✅	3m 19s
k8s_test_crs_uniqueness	✅	33.0s
k8s_test_insecure_keygen_encrypt_and_public_decrypt	✅	3m 17s
k8s_test_insecure_keygen_encrypt_multiple_types	✅	3m 37s
k8s_test_keygen_and_crs	✅	3m 14s
k8s_test_keygen_uniqueness	✅	8m 52s
nightly_full_gen_tests_k8s_default_centralzd_sequential_crs	✅	1.8s
test_k8s_centralzd_insecure	✅	5m 2s
k8s_test_centralized_insecure	✅	1m
nightly_full_gen_tests_default_k8s_centralized_sequential_crs	✅	1.7s
nightly_full_gen_tests_k8s_default_centralzd_sequential_crs	✅	1.7s
test_k8s_centralzd_insecure	✅	1m 2s
k8s_test_centralized_insecure	✅	1m 1s
nightly_full_gen_tests_default_k8s_centralized_sequential_crs	✅	1.7s

_{🍂 No flaky tests in this run.}

Github Test Reporter by CTRF 💚

🔄 This comment has been updated

github-actions · 2026-03-05T14:47:18Z

Vulnerability Scan Results

Details


Report Summary

┌───────────────────────────────────┬────────────┬─────────────────┬─────────┐
│              Target               │    Type    │ Vulnerabilities │ Secrets │
├───────────────────────────────────┼────────────┼─────────────────┼─────────┤
│ base:latest (chainguard 20230214) │ chainguard │        0        │    -    │
├───────────────────────────────────┼────────────┼─────────────────┼─────────┤
│ usr/bin/yq                        │  gobinary  │        3        │    -    │
└───────────────────────────────────┴────────────┴─────────────────┴─────────┘
Legend:
- '-': Not scanned
- '0': Clean (no security findings detected)


For OSS Maintainers: VEX Notice
--------------------------------
If you're an OSS maintainer and Trivy has detected vulnerabilities in your project that you believe are not actually exploitable, consider issuing a VEX (Vulnerability Exploitability eXchange) statement.
VEX allows you to communicate the actual status of vulnerabilities in your project, improving security transparency and reducing false positives for your users.
Learn more and start using VEX: https://trivy.dev/docs/v0.69/guide/supply-chain/vex/repo#publishing-vex-documents

To disable this notice, set the TRIVY_DISABLE_VEX_NOTICE environment variable.


usr/bin/yq (gobinary)
=====================
Total: 3 (HIGH: 2, CRITICAL: 1)

┌─────────┬────────────────┬──────────┬────────┬───────────────────┬──────────────────────────────┬──────────────────────────────────────────────────────────────┐
│ Library │ Vulnerability  │ Severity │ Status │ Installed Version │        Fixed Version         │                            Title                             │
├─────────┼────────────────┼──────────┼────────┼───────────────────┼──────────────────────────────┼──────────────────────────────────────────────────────────────┤
│ stdlib  │ CVE-2025-68121 │ CRITICAL │ fixed  │ v1.25.5           │ 1.24.13, 1.25.7, 1.26.0-rc.3 │ crypto/tls: Unexpected session resumption in crypto/tls      │
│         │                │          │        │                   │                              │ https://avd.aquasec.com/nvd/cve-2025-68121                   │
│         ├────────────────┼──────────┤        │                   ├──────────────────────────────┼──────────────────────────────────────────────────────────────┤
│         │ CVE-2025-61726 │ HIGH     │        │                   │ 1.24.12, 1.25.6              │ golang: net/url: Memory exhaustion in query parameter        │
│         │                │          │        │                   │                              │ parsing in net/url                                           │
│         │                │          │        │                   │                              │ https://avd.aquasec.com/nvd/cve-2025-61726                   │
│         ├────────────────┤          │        │                   │                              ├──────────────────────────────────────────────────────────────┤
│         │ CVE-2025-61728 │          │        │                   │                              │ golang: archive/zip: Excessive CPU consumption when building │
│         │                │          │        │                   │                              │ archive index in archive/zip                                 │
│         │                │          │        │                   │                              │ https://avd.aquasec.com/nvd/cve-2025-61728                   │
└─────────┴────────────────┴──────────┴────────┴───────────────────┴──────────────────────────────┴──────────────────────────────────────────────────────────────┘

dvdplm

LGMT

core/service/src/heap_profiling.rs

dvdplm · 2026-03-05T15:34:00Z

profiling/analyze-heap.sh

+    fi
+}
+
+# sed -i: GNU sed uses -i '', BSD (macOS) sed requires -i ''


This seems to imply that they are the same?

Fixed in 44b0f8a (they are not the same apparently, but doc was unclear indeed)

core/service/Cargo.toml

profiling/README.md

kc1212 · 2026-03-05T16:14:22Z

I've been trying this PR on mac OS but I can't figure out how to run pprof. The script profiling/analyze-heap.sh tells me to install gperftools using homebrew but when I install that there doesn't seem to be an executable.

But this is not super surprising since I've always had problems with profilers on mac OS.

If anyone has ideas please let me know, otherwise I can try compiling gperftools from scratch and see if it gives me a pporf executable.

kc1212 · 2026-03-05T16:24:20Z

Ok I installed pprof using golang https://github.com/google/pprof

but then I got the error

Found 1 heap dump(s)
./profiling/analyze-heap.sh: line 170: DUMPS_ORIG: bad array subscript

It seems not all bash implementation supports negative indexing

eudelins-zama · 2026-03-05T16:31:02Z

@kc1212 pprof != jeprof AFAIU

eudelins-zama · 2026-03-05T16:32:42Z

Btw, feel free to push an update on macOS setup documentation if you manage to troubleshoot some issues on your side!

kc1212 · 2026-03-05T16:32:51Z

@kc1212 pprof != jeprof AFAIU

well, the script is asking for jeprof, google-pprof or pprof

eudelins-zama · 2026-03-05T16:36:19Z

Docs says:

## Host dependencies

- jeprof (from gperftools) — reads jemalloc .heap dumps
- graphviz — renders SVG flamegraphs (
- addr2line (from binutils) — resolves addresses to source lines

But indeed the script seems to accept pprof which it should not

profiling/analyze-heap.sh

kc1212 · 2026-03-05T16:42:22Z

Btw, feel free to push an update on macOS setup documentation if you manage to troubleshoot some issues on your side!

pushed the changes, seems to be all ok now!

core/service/config/compose_1.toml

titouantanguy · 2026-03-10T16:06:26Z

Small note that you configure jemalloc to auto-dump to /tmp/kms-heap/auto but this folder doesn't exist by default in the docker container we use for profiling (leading to jemalloc complaining).

eudelins-zama added 2 commits March 3, 2026 11:37

fix: add missing field in config

f5ff202

feat: add jemalloc heap profiling infrastructure

e663733

cla-bot bot added the cla-signed The CLA has been signed. label Mar 5, 2026

chore: simplify analysis script

2ccf39c

fix: trivy action bump

ba79db2

eudelins-zama self-assigned this Mar 5, 2026

eudelins-zama marked this pull request as ready for review March 5, 2026 14:58

eudelins-zama requested a review from a team as a code owner March 5, 2026 14:58

dvdplm approved these changes Mar 5, 2026

View reviewed changes

eudelins-zama added 4 commits March 5, 2026 16:57

chore: rename jemalloc-stats to heap-profiling

298351a

fix: use std::ffi::CString

b5b2623

docs: portable sed doc clarification

44b0f8a

docs: memory profiling doc clarification

16386f0

eudelins-zama commented Mar 5, 2026

View reviewed changes

profiling/analyze-heap.sh Outdated Show resolved Hide resolved

chore: fix profiler script for macOS

2b69f54

dd23 reviewed Mar 5, 2026

View reviewed changes

core/service/config/compose_1.toml Outdated Show resolved Hide resolved

eudelins-zama added 2 commits March 5, 2026 17:58

fix: revert config update

fa1e321

fix: heap dump diff analysis fix

f971458

dd23 mentioned this pull request Mar 9, 2026

chore: bump rust, yq, grpc-health-check, tfhe-rs and trivy #455

Merged

15 tasks

dd23 closed this in #455 Mar 10, 2026

dd23 reopened this Mar 10, 2026

Conversation

eudelins-zama commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

What's included

Issue ticket number and link

PR Checklist

Dependency Update Questionnaire (only if deps changed or added)

Uh oh!

eudelins-zama commented Mar 5, 2026

Usage

Uh oh!

github-actions bot commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Consolidated Tests Results 2026-03-08 - 20:15:43

Test Results

Details

test-reporter: Run #749

🎉 All tests passed!

Tests

Uh oh!

github-actions bot commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Vulnerability Scan Results

Uh oh!

dvdplm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dvdplm Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

eudelins-zama Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kc1212 commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kc1212 commented Mar 5, 2026

Uh oh!

eudelins-zama commented Mar 5, 2026

Uh oh!

eudelins-zama commented Mar 5, 2026

Uh oh!

kc1212 commented Mar 5, 2026

Uh oh!

eudelins-zama commented Mar 5, 2026

Uh oh!

Uh oh!

kc1212 commented Mar 5, 2026

Uh oh!

Uh oh!

titouantanguy commented Mar 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

eudelins-zama commented Mar 5, 2026 •

edited

Loading

github-actions bot commented Mar 5, 2026 •

edited

Loading

github-actions bot commented Mar 5, 2026 •

edited

Loading

kc1212 commented Mar 5, 2026 •

edited

Loading